// Copyright Epic Games, Inc. All Rights Reserved.

#pragma once

#include <zencore/timer.h>
#include <zenstore/hashkeyset.h>

namespace zen {

class WorkerThreadPool;

/** Context object for data scrubbing

  Data scrubbing is when we traverse stored data to validate it and
  optionally correct/recover
 */

class ScrubContext
{
public:
	ScrubContext(WorkerThreadPool&					   InWorkerThreadPool,
				 std::chrono::steady_clock::time_point Deadline = std::chrono::steady_clock::time_point::max());
	~ScrubContext();

	void			ReportBadCidChunks(std::span<IoHash> BadCasChunks);
	inline uint64_t ScrubTimestamp() const { return m_ScrubTime; }
	void			ReportScrubbed(uint64_t ChunkCount, uint64_t ChunkBytes)
	{
		m_ChunkCount.fetch_add(ChunkCount);
		m_ByteCount.fetch_add(ChunkBytes);
	}

	std::chrono::steady_clock::time_point GetDeadline() const { return m_Deadline; }
	bool								  IsWithinDeadline() const;
	void								  ThrowIfDeadlineExpired() const;

	inline uint64_t ScrubbedChunks() const { return m_ChunkCount; }
	inline uint64_t ScrubbedBytes() const { return m_ByteCount; }

	HashKeySet BadCids() const;
	bool	   IsBadCid(const IoHash& Cid) const;

	inline bool RunRecovery() const { return m_Recover; }
	inline void SetShouldDelete(bool DoDelete) { m_Recover = DoDelete; }

	inline bool IsSkipCas() const { return m_SkipCas; }
	inline void SetSkipCas(bool DoSkipCas) { m_SkipCas = DoSkipCas; }

	inline WorkerThreadPool& ThreadPool() { return m_WorkerThreadPool; }

private:
	uint64_t							  m_ScrubTime = GetHifreqTimerValue();
	bool								  m_Recover	  = true;
	bool								  m_SkipCas	  = false;
	std::atomic<uint64_t>				  m_ChunkCount{0};
	std::atomic<uint64_t>				  m_ByteCount{0};
	mutable RwLock						  m_Lock;
	HashKeySet							  m_BadCid;
	WorkerThreadPool&					  m_WorkerThreadPool;
	std::chrono::steady_clock::time_point m_Deadline{};
};

class ScrubDeadlineExpiredException : public std::runtime_error
{
public:
	ScrubDeadlineExpiredException();
	~ScrubDeadlineExpiredException();
};

}  // namespace zen
